"""
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    created by lane.chang on '28/10/2024'
    comment: 优必选
"""
import asyncio
import aiohttp
from project.lib.decorator import exec_duration

CONTEXT_TOKEN_LIMIT = 3000

class Ubtech:

    @staticmethod
    @exec_duration()
    async def post(user_text, system_text='', model='qwen', stream=False, max_tokens=None, temperature=None):
        """
        :param user_text:
        :param system_text:
        :param model:
        :param stream:
        :param max_tokens:
        :param temperature:
        :return:
        """
        base_url = 'http://10.20.223.83:62474/v1/chat/completions'
        headers = {
            'Content-Type': 'application/json'
        }

        params = dict()
        params['model'] = model
        messages = []
        if system_text:
            messages.append({'role': 'system', 'content': system_text})
        if user_text:
            messages.append({'role': 'user', 'content': user_text})
        params['messages'] = messages
        if max_tokens is not None:
            params['max_tokens'] = max_tokens
        if temperature is not None:
            params['temperature'] = temperature
        params['stream'] = stream

        async with aiohttp.ClientSession() as session:
            async with session.post(base_url, headers=headers, json=params) as resp:
                if resp.status != 200:
                    raise Exception(f'研究院模型推理失败 status: {resp.status} url: {base_url}')

                result = await resp.json()

        return result['choices'][0]['message']['content']

    @staticmethod
    @exec_duration()
    async def get_embedding(text, model='bge-large-zh-v1.5') -> list:
        """ 生成向量数据
        :param text:
        :param model:
        :return:
        """
        base_url = 'http://10.20.223.89:9997/v1/embeddings'
        headers = {
            'Content-Type': 'application/json'
        }

        params = {
            'model': model,
            'input': text
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(base_url, headers=headers, json=params) as resp:
                if resp.status != 200:
                    raise Exception(f'研究院模型embedding失败 status: {resp.status} url: {base_url}')

                result = await resp.json()

        return result['data'][0]['embedding']

@exec_duration()
async def main():
    """
    :return:
    """
    tasks = list()
    user_message1 = '给我列举五条狗狗的名称'
    user_message2 = '给我简单介绍一下中国'
    user_message3 = '给我简单介绍一下唐朝'
    user_message4 = '给我简单介绍一下宋朝'
    user_message5 = '给我简单介绍一下元朝'

    tasks.append(
        asyncio.create_task(Ubtech.post(user_text=user_message1), name='task1')
    )

    tasks.append(
        asyncio.create_task(Ubtech.post(user_text=user_message2), name='task2')
    )

    tasks.append(
        asyncio.create_task(Ubtech.post(user_text=user_message3), name='task3')
    )

    tasks.append(
        asyncio.create_task(Ubtech.post(user_text=user_message4), name='task4')
    )

    tasks.append(
        asyncio.create_task(Ubtech.post(user_text=user_message5), name='task5')
    )

    await asyncio.gather(*tasks)

    for task in tasks:
        print(task.result())


if __name__ == '__main__':

    result = asyncio.run(Ubtech.get_embedding('我是畅垒'))
    print(result)

